import torch
import torch.nn as nn
import torch.nn.functional as F

class TextRNN(nn.Module):
    def __init__(self, embed_dim=200, hidden_size=128, num_layers=2, class_num=4):
        super().__init__()
        self.lstm = nn.LSTM(embed_dim, hidden_size, num_layers, bidirectional=True, batch_first=True, dropout=0.5)
        self.fc = nn.Linear(hidden_size * 2, class_num)
    
    def forward(self, x):
        """
        Compute the forward pass of TextRNN.
        
        Args:
        - x: (batch, len, w2v_dim) torch.
        
        Returns:
        - logit: (batch, 4) the corresponding logits.
        """
        out, _ = self.lstm(out)
        out = self.fc(out[:, -1, :])  # 句子最后时刻的 hidden state
        logit = F.log_softmax(x, dim=1)
        return logit
    
    def loss(self):
        """
        Compute the loss.
        
        Args:
        - y: :torch.tensor: (batch,) the prediction.
        - labels: :torch.LongTensor: (batch,) the true labels.
        
        Returns:
        - loss: :torch.tensor: the corresponding tensor.
        """   
        loss = F.nll_loss(y, labels)
        return loss